#include <cusp/array2d.h>
#include <cusp/io/matrix_market.h>
#include <cusp/multiply.h>
#include <cusp/device/block/qr.h>
#include <cusp/qr.h>

#include <iostream>

#include "../timer.h"

/** Check one matrix. */
template<typename ReferenceMatrix,
	 typename DeviceMatrix1,
	 typename DeviceMatrix2>
int check_result(ReferenceMatrix& Ahost, const DeviceMatrix1& Q, const DeviceMatrix2& R)
{
	typedef ReferenceMatrix HostMatrix;
	int i = 0, f = 0;

	HostMatrix Qhost = Q;
	HostMatrix Rhost = R;
	HostMatrix Bhost;

	cusp::multiply(Qhost, Rhost, Bhost);

	if (Ahost.num_rows != Bhost.num_rows || Ahost.num_cols != Bhost.num_cols)
	{
		std::cerr << "Check FAILED! (size)" << std::endl;
		return -1;
	}

	typedef typename HostMatrix::values_array_type::iterator Iterator;
	Iterator first1 = Ahost.values.begin();
	Iterator last1  = Ahost.values.end();
	Iterator first2 = Bhost.values.begin();
	while(first1 != last1)
	{
		/* same condition used in unit testing */
		if (abs(*first1 - *first2) > 1e-4 * (abs(*first1) + abs(*first2)) + 1e-4)
		{
			++f;
		}
		first1++;
		first2++;
		++i;
	}

	if (f != 0)
		std::cerr << "Check FAILED! (" << f << ")" << std::endl;

	return f;
}

/** Check multiple matrices in contiguous memory. */
template<typename ReferenceMatrix,
	 typename RandomAccessIterator1,
	 typename RandomAccessIterator2>
int check_result(ReferenceMatrix& A,
		 size_t num_rows, size_t num_cols,
		 RandomAccessIterator1 Qs,
		 RandomAccessIterator2 Rs,
		 size_t nMat)
{
	int result = 0;
	for (size_t i = 0; i < nMat; ++i)
	{
		result += check_result(A,
			cusp::make_array2d_view(num_rows, num_rows, num_rows,
				cusp::make_array1d_view(Qs, Qs + (num_rows*num_rows)),
				cusp::column_major()),
			cusp::make_array2d_view(num_rows, num_cols, num_cols,
				cusp::make_array1d_view(Rs, Rs + (num_rows*num_cols)),
				cusp::row_major())
			);
		Qs += num_rows*num_rows;
		Rs += num_rows*num_cols;
	}
	return result;
}


template<typename RandomAccessIterator1,
	 typename RandomAccessIterator2>
__global__
void block_qr_kernel(RandomAccessIterator1 As,
		     size_t num_rows,
		     size_t num_cols,
		     RandomAccessIterator2 Qs,
		     size_t nMat)
{
	for (size_t i = blockIdx.x; i < nMat; i += gridDim.x)
	{
		/* handle i-th matrix */
		cusp::device::block::qr(
			As + (i * num_rows * num_cols),
			num_rows,
			num_cols,
			Qs + (i * num_rows * num_rows));
	}
}


/** Read one matrix from a file and do QR decomposition multiple times,
    possibly in parallel.

    We always perform the decomposition of the same matrix, in order
    to keep the program simple.

    Usage: ./blockqr <matrix_file.mtx> <nMat>

    The returned time is the total for performing nMat times
    the decomposition, in seconds.
*/
int main(int argc, char** argv)
{
	typedef int    IndexType;
	typedef float  ValueType;

	typedef cusp::array2d<ValueType,cusp::host_memory>   HostMatrix;
	typedef cusp::array2d<ValueType,cusp::device_memory> DeviceMatrix;

	if (argc < 2)
	{
		std::cout << "Missing argument." << std::endl;
		return 1;
	}

	/* number of matrices */
	size_t nMat = 0;
	if (argc > 2)
	{
		nMat = atoi(argv[2]);
	}
	if (nMat < 1)
	{
		std::cout << "Too few repetitions." << std::endl;
		return 1;
	}

	/* get the matrix from file */
	HostMatrix A;
	std::cerr << "Reading matrix from file: " << argv[1] << std::endl;
	cusp::io::read_matrix_market_file(A, std::string(argv[1]));
	std::cerr << "size " << A.num_rows << " x " << A.num_cols
		<< "\ndoing " << nMat << " decompositions" << std::endl;


	/* Use block::qr */
	{

	/* set up matrices in device memory */
	size_t size = A.num_rows * A.num_cols;
	cusp::array1d<ValueType,cusp::device_memory> As(size * nMat);
	cusp::array1d<ValueType,cusp::device_memory> Qs(A.num_rows * A.num_rows * nMat, ValueType(0));
	cusp::array1d<ValueType,cusp::device_memory>::iterator first = As.begin();
	cusp::array1d<ValueType,cusp::device_memory>::iterator firstEnd = first + size;
	cusp::array1d<ValueType,cusp::device_memory>::iterator a = firstEnd;
	thrust::copy(A.values.begin(), A.values.end(), first);
	for (size_t i = 1; i < nMat; ++i)
	{
		thrust::copy(first, firstEnd, a);
		a += size;
	}

	/* timing */
	timer t;
	block_qr_kernel<<<thrust::min<size_t>(nMat,1000),128>>>(As.begin(), A.num_rows, A.num_cols, Qs.begin(), nMat);
	float time = t.seconds_elapsed();
	std::cout << time << "\t";
	std::cout.flush();

	/* check result */
	check_result(A, A.num_rows, A.num_cols, Qs.begin(), As.begin(), nMat);

	} /* end block::qr */


	/* Use cusp::qr */
	{

	/* set up matrices */
	/* always reuse the same ones */
	DeviceMatrix Adev = A;
	DeviceMatrix Qdev, Rdev;

	/* timing */
	timer t;
	for (size_t i = 0; i < nMat; ++i)
	{
		cusp::qr(Adev, Qdev, Rdev);
	}
	float time = t.seconds_elapsed();
	std::cout << time << std::endl;

	/* check result (only the last one) */
	check_result(A, Qdev, Rdev);

	} /* end cusp::qr */

	return 0;
}

